import datetime
import uuid
from typing import List, Dict

from fastapi import Request
from langchain_core.documents import Document
from pydantic import BaseModel
from loguru import logger

from configs.kb_config import ES_KB_INDEX_NAME, ES_KB_EMERGENCY_NAME
from db.repository.qa_file import db_get_qa_file_by_id
from schemas.response_entity import JsonModel
from service.es_service import ELASTICSEARCH, ELASTICSEARCH_EMERGENCY
from service.embeddings import embeddings
# from service.es_service.es_operate import es_update_chunk_by_id, es_delete_chunk_by_id


class UpdateChunkEntity(BaseModel):
	token: str
	file_id: uuid.UUID
	chunk_id: str
	content: str = None
	standard_name: str = None
	standard_number: str = None
	standard_clause: str = None
	standard_tags: List[str] = None


class DeleteChunkEntity(BaseModel):
	token: str
	chunk_id: str
	file_id: uuid.UUID


async def show_file_chunks(request: Request, file_id: uuid.UUID):
	"""
	获取文件的所有chunks
	"""
	user = request.state.user

	if not user.check_admin_role():
		return JsonModel(code=401, data=None, msg="用户没有权限").to_response()

	file = db_get_qa_file_by_id(file_id)
	if not file:
		return JsonModel(code=400, data=None, msg="文件不存在").to_response()

	query_body = {
		"_source": ["text", "metadata"],
		"query": {
			"term": {
				"metadata.file_id.keyword": str(file_id)
			}
		},
		"size": 1000
	}

	search_result = []
	try:
		# 搜索主索引
		result1 = ELASTICSEARCH.search(idxnm=ES_KB_INDEX_NAME, q=query_body)
		if result1 and 'hits' in result1:
			search_result.extend(result1['hits']['hits'])

		# 搜索应急索引
		result2 = ELASTICSEARCH_EMERGENCY.search(idxnm=ES_KB_EMERGENCY_NAME, q=query_body)
		if result2 and 'hits' in result2:
			search_result.extend(result2['hits']['hits'])
	except Exception as e:
		logger.error(f"ES search error: {e}")
		return JsonModel(code=500, data=None, msg="搜索失败").to_response()

	documents = []
	for hit in search_result:
		doc = hit['_source']
		metadata = {k: v for k, v in doc.get("metadata", {}).items() if k in ['file_id', 'file_name']}
		document = Document(id=hit["_id"], page_content=doc.get("text"), metadata=metadata)
		documents.append(document)
	return JsonModel(code=200, data=documents, msg="Success").to_response()


async def update_file_chunk(request: Request, payload: UpdateChunkEntity):
	"""
	更新某一条chunk，应该是要用chunk_id
	"""
	user = request.state.user

	if not user.check_admin_role():
		return JsonModel(code=401, data=None, msg="用户没有权限").to_response()

	file = db_get_qa_file_by_id(payload.file_id)
	if not file:
		return JsonModel(code=400, data=None, msg="文件不存在").to_response()

	# 先尝试在主索引中查找
	es_doc = ELASTICSEARCH.get(idxnm=ES_KB_INDEX_NAME, doc_id=payload.chunk_id, error_trace=False, ignore=[404])
	found = es_doc.get("found", False)
	index_name = ES_KB_INDEX_NAME
	
	# 如果主索引中没找到，尝试在应急索引中查找
	if not found:
		es_doc = ELASTICSEARCH_EMERGENCY.get(idxnm=ES_KB_EMERGENCY_NAME, doc_id=payload.chunk_id, error_trace=False, ignore=[404])
		found = es_doc.get("found", False)
		index_name = ES_KB_EMERGENCY_NAME
		es_client = ELASTICSEARCH_EMERGENCY
	else:
		es_client = ELASTICSEARCH
	
	if not found:
		return JsonModel(code=400, data=None, msg=f"chunk不存在").to_response()
	
	# 获取文档内容
	doc = es_doc.get("_source", {}).copy()
	metadata = doc.get("metadata", {}).copy()
	
	# 更新文本内容（如果提供了）
	if payload.content is not None:
		doc["text"] = payload.content
		doc["vector"] = embeddings.embed_query(payload.content)
	
	# 更新metadata中的标准相关字段
	if payload.standard_name is not None:
		metadata["standard_name"] = payload.standard_name
	if payload.standard_number is not None:
		metadata["standard_number"] = payload.standard_number
	if payload.standard_clause is not None:
		metadata["standard_clause"] = payload.standard_clause
	if payload.standard_tags is not None:
		metadata["standard_tags"] = payload.standard_tags
	
	doc["metadata"] = metadata
	
	flag = es_client.update(doc_id=payload.chunk_id, body={"doc": doc}, idxnm=index_name)
	if flag is True:
		return JsonModel(code=200, data=None, msg=f"更新成功").to_response()
	else:
		return JsonModel(code=500, data=None, msg=f"更新失败, error{flag}").to_response()


async def delete_file_chunk(request: Request, payload: DeleteChunkEntity):
	"""
	删除某一条chunk
	"""
	user = request.state.user
	if user is None:
		return JsonModel(code=401, data=None, msg="用户token非法").to_response()

	if not user.check_admin_role():
		return JsonModel(code=401, data=None, msg="用户没有权限").to_response()

	file = db_get_qa_file_by_id(payload.file_id)
	if not file:
		return JsonModel(code=400, data=None, msg="文件不存在").to_response()

	# 先尝试在主索引中查找
	es_doc = ELASTICSEARCH.get(idxnm=ES_KB_INDEX_NAME, doc_id=payload.chunk_id, error_trace=False, ignore=[404])
	found = es_doc.get("found", False)
	index_name = ES_KB_INDEX_NAME
	
	# 如果主索引中没找到，尝试在应急索引中查找
	if not found:
		es_doc = ELASTICSEARCH_EMERGENCY.get(idxnm=ES_KB_EMERGENCY_NAME, doc_id=payload.chunk_id, error_trace=False, ignore=[404])
		found = es_doc.get("found", False)
		index_name = ES_KB_EMERGENCY_NAME
		es_client = ELASTICSEARCH_EMERGENCY
	else:
		es_client = ELASTICSEARCH
	
	if not found:
		return JsonModel(code=400, data=None, msg=f"chunk不存在").to_response()
	
	flag = es_client.delete(doc_id=payload.chunk_id, idxnm=index_name)
	if flag is True:
		return JsonModel(code=200, data=None, msg=f"删除成功").to_response()
	else:
		return JsonModel(code=500, data=None, msg=f"删除失败, error{flag}").to_response()

